#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
from ..._main import PamGroup, PamControl, PamRule, FeatureBase
from ..._error import VerifyError, ModifyError, FeatureNotFoundError
from ..._feature_common import Fc


class RootOk(FeatureBase):

    SO = "pam_rootok.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        if Fc.groupsHasModule(self._pamService._data, (PamGroup.AUTH, False), self.SO):
            raise VerifyError("%s can only be used for group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

        ruleList, indexList = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        if len(ruleList) == 0:
            raise FeatureNotFoundError()
        if len(ruleList) > 1:
            raise VerifyError("%s is used multiple times" % (self.SO), self._pamService.file_path)
        if ruleList[0].module_arguments is not None:
            raise VerifyError("%s should not have any argument" % (self.SO), self._pamService.file_path)
        if ruleList[0].control != PamControl.SUFFICIENT:
            raise VerifyError("%s can only be used for control \"%s\"" % (self.SO, PamControl.SUFFICIENT), self._pamService.file_path)
        if indexList[0] != 0:
            raise VerifyError("%s must be the first rule in group" % (self.SO), self._pamService.file_path)

    @staticmethod
    def name():
        return "root-ok"

    def verify_after_init(self):
        pass

    def get_rule_index(self):
        _, indexList = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        assert len(indexList) == 1
        return indexList[0]

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            # pam_rootok.so should always be the first rule
            Fc.putPamGroup(target_data, PamGroup.AUTH).insert(0, PamRule(PamControl.SUFFICIENT, self.SO))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class Shells(FeatureBase):

    SO = "pam_shells.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        shellsFile = "/etc/shells"

        if not os.path.exists(os.path.join(context.dir_prefix, shellsFile[1:])):
            if Fc.groupsHasModule(self._pamService._data, "*", self.SO):
                raise VerifyError("%s used when %s does not exist" % (self.SO, shellsFile), self._pamService.file_path)
            raise FeatureNotFoundError()

        if Fc.groupsHasModule(self._pamService._data, (PamGroup.AUTH, False), self.SO):
            raise VerifyError("%s can only be used for group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

        ruleList, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        if len(ruleList) == 0:
            raise FeatureNotFoundError()
        if len(ruleList) > 1:
            raise VerifyError("%s used for multiple times" % (self.SO), self._pamService.file_path)
        if ruleList[0].module_arguments is not None:
            raise VerifyError("%s should not have any argument" % (self.SO), self._pamService.file_path)
        if ruleList[0].control != PamControl.REQUIRED:
            raise VerifyError("%s can only be used for control \"%s\"" % (self.SO, PamControl.REQUIRED), self._pamService.file_path)

    @staticmethod
    def name():
        return "shells"

    def verify_after_init(self):
        # must after pam_rootok.so, root can always login
        if not Fc.isAfterInGroupIfExist(self._pamService, self, RootOk):
            raise VerifyError("%s must be after %s in group" % (self.SO, RootOk.SO), self._pamService.file_path)

    def get_rule_index(self):
        _, indexList = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        assert len(indexList) == 1
        return indexList[0]

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            afterIndex = Fc.changeAfterIndexIfExist(self._pamService, RootOk, -1)       # after pam_rootok.so, root can always login
            Fc.putPamGroup(target_data, PamGroup.AUTH).insert(afterIndex + 1, PamRule())        # FIXME
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class NoLogin(FeatureBase):

    SO = "pam_nologin.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        ruleListAuth, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        ruleListAcct, _ = Fc.getPamGroup(self._pamService._data, PamGroup.ACCOUNT, lambda x: x.module_path == self.SO)
        if len(ruleListAuth) == 0 and len(ruleListAcct) == 0:
            raise FeatureNotFoundError()
        if (len(ruleListAuth) > 0 and len(ruleListAcct) == 0) or (len(ruleListAuth) == 0 and len(ruleListAcct) > 0):
            raise VerifyError("%s must be used for group \"%s\" and \"%s\" simultaneously" % (self.SO, PamGroup.AUTH, PamGroup.ACCOUNT), self._pamService.file_path)

        for pamGroup, ruleList in {PamGroup.AUTH: ruleListAuth, PamGroup.ACCOUNT: ruleListAcct}:
            if len(ruleList) > 1:
                raise VerifyError("%s used for multiple times in group \"%s\"" % (self.SO, pamGroup), self._pamService.file_path)
            if ruleList[0].module_arguments is not None:
                raise VerifyError("%s in group \"%s\" should not have any argument" % (self.SO, pamGroup), self._pamService.file_path)
            if ruleList[0].control != PamControl.REQUIRED:
                raise VerifyError("%s in group \"%s\" can only be used for control \"%s\"" % (self.SO, pamGroup, PamControl.REQUIRED), self._pamService.file_path)

    @staticmethod
    def name():
        return "no-login"

    def verify_after_init(self):
        # must after pam_rootok.so, root can always login
        if not Fc.isAfterInGroupIfExist(self._pamService, self, RootOk):
            raise VerifyError("%s must be after %s in group" % (self.SO, RootOk.SO), self._pamService.file_path)

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            # group auth
            if True:
                afterIndex = Fc.changeAfterIndexIfExist(self._pamService, RootOk, -1)       # after pam_rootok.so, root can always login
                Fc.putPamGroup(target_data, PamGroup.AUTH).insert(afterIndex + 1, PamRule(PamControl.REQUIRED, self.SO))
            # group account
            if True:
                Fc.putPamGroup(target_data, PamGroup.ACCOUNT).insert(0, PamRule(PamControl.REQUIRED, self.SO))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class WheelOk(FeatureBase):

    SO = "pam_wheel.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        if Fc.groupsHasModule(self._pamService._data, (PamGroup.AUTH, False), self.SO):
            raise VerifyError("%s can only be used for group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

        ruleList, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        if len(ruleList) == 0:
            raise FeatureNotFoundError()
        if len(ruleList) > 1:
            raise VerifyError("%s used for multiple times" % (self.SO), self._pamService.file_path)
        if ruleList[0].control != PamControl.SUFFICIENT:
            raise FeatureNotFoundError()                                     # pam_wheel.so is used, but it is not for "wheel ok"
        if not Fc.moduleArgumentsEquals(ruleList[0].module_arguments, set(["trust", "use_uid"])):
            raise FeatureNotFoundError()                                     # pam_wheel.so is used, but it is not for "wheel ok"

    @staticmethod
    def name():
        return "wheel-ok"

    def verify_after_init(self):
        # must after pam_rootok.so, wheel should be lower than root
        if not Fc.isAfterInGroupIfExist(self._pamService, self, RootOk):
            raise VerifyError("%s must be after %s in group" % (self.SO, RootOk.SO), self._pamService.file_path)
        # after pam_shell.so, wheel should also obey shell restriction
        if not Fc.isAfterInGroupIfExist(self._pamService, self, Shells):
            raise VerifyError("%s must be after %s in group" % (self.SO, Shells.SO), self._pamService.file_path)

    def get_rule_index(self):
        _, indexList = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        assert len(indexList) == 1
        return indexList[0]

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            if Fc.groupsHasModule(target_data, "*", self.SO):
                raise ModifyError("%s exists" % (self.SO), self._pamService.file_path)
            afterIndex = -1
            afterIndex = Fc.changeAfterIndexIfExist(self._pamService, RootOk, afterIndex)       # after pam_rootok.so, wheel should be lower than root
            afterIndex = Fc.changeAfterIndexIfExist(self._pamService, Shells, afterIndex)       # after pam_shell.so, wheel should also obey shell restriction
            Fc.putPamGroup(target_data, PamGroup.AUTH).insert(afterIndex + 1, PamRule(PamControl.SUFFICIENT, self.SO, ["trust", "use_uid"]))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class MustBeWheel(FeatureBase):

    SO = "pam_wheel.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        if Fc.groupsHasModule(self._pamService._data, (PamGroup.AUTH, False), self.SO):
            raise VerifyError("%s can only be used for group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

        ruleList, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        if len(ruleList) == 0:
            raise FeatureNotFoundError()
        if len(ruleList) > 1:
            raise VerifyError("%s used for multiple times" % (self.SO), self._pamService.file_path)
        if ruleList[0].control != PamControl.REQUIRED:
            raise FeatureNotFoundError()                                     # pam_wheel.so is used, but it is not for "must be wheel"
        if Fc.moduleArgumentsContains(ruleList[0].module_arguments, ["trust"]):
            raise FeatureNotFoundError()                                     # pam_wheel.so is used, but it is not for "must be wheel"

    @staticmethod
    def name():
        return "must-in-wheel"

    def verify_after_init(self):
        # must after pam_rootok.so, wheel should be lower than root
        if not Fc.isAfterInGroupIfExist(self._pamService, self, RootOk):
            raise VerifyError("%s must be after %s in group" % (self.SO, RootOk.SO), self._pamService.file_path)
        # after pam_shell.so, wheel should also obey shell restriction
        if not Fc.isAfterInGroupIfExist(self._pamService, self, Shells):
            raise VerifyError("%s must be after %s in group" % (self.SO, Shells.SO), self._pamService.file_path)

    def get_rule_index(self):
        _, indexList = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        assert len(indexList) == 1
        return indexList[0]

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            if Fc.groupsHasModule(target_data, "*", self.SO):
                raise ModifyError("%s exists" % (self.SO), self._pamService.file_path)
            afterIndex = -1
            afterIndex = Fc.changeAfterIndexIfExist(self._pamService, RootOk, afterIndex)       # after pam_rootok.so, wheel should be lower than root
            afterIndex = Fc.changeAfterIndexIfExist(self._pamService, Shells, afterIndex)       # after pam_shell.so, wheel should also obey shell restriction
            Fc.putPamGroup(target_data, PamGroup.AUTH).insert(afterIndex + 1, PamRule(PamControl.REQUIRED, self.SO, ["use_uid"]))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class UnixLogin(FeatureBase):

    SO = "pam_unix.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        ruleListAuth, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        ruleListAcct, _ = Fc.getPamGroup(self._pamService._data, PamGroup.ACCOUNT, lambda x: x.module_path == self.SO)
        if len(ruleListAuth) == 0 and len(ruleListAcct) == 0:
            raise FeatureNotFoundError()
        if (len(ruleListAuth) > 0 and len(ruleListAcct) == 0) or (len(ruleListAuth) == 0 and len(ruleListAcct) > 0):
            raise VerifyError("%s must be used for group \"%s\" and \"%s\" simultaneously" % (self.SO, PamGroup.AUTH, PamGroup.ACCOUNT), self._pamService.file_path)

        # pam_unix.so in pam group "auth"
        pass

        # pam_unix.so in pam group "account"
        if len(ruleListAcct) > 1:
            raise VerifyError("%s used for multiple times in group \"%s\"" % (self.SO, PamGroup.ACCOUNT), self._pamService.file_path)
        if ruleListAcct[0].module_arguments is not None:
            raise VerifyError("%s in group \"%s\" should not have any argument" % (self.SO, PamGroup.ACCOUNT), self._pamService.file_path)
        if ruleListAcct[0].control != PamControl.REQUIRED:
            raise VerifyError("%s in group \"%s\" can only be used for control \"%s\"" % (self.SO, PamGroup.ACCOUNT, PamControl.REQUIRED), self._pamService.file_path)

    @staticmethod
    def name():
        return "unix-login"

    def verify_after_init(self):
        pass

    def modify(self, target_data, add_or_remove):
        assert False


class _DenyOrPermit(FeatureBase):

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        for group in [PamGroup.AUTH, PamGroup.ACCOUNT]:
            rules = Fc.getPamGroup(self._pamService._data, group)[0]
            if len(rules) != 1 or rules[0].module_path != self.SO or rules[0].control != PamControl.REQUIRED:
                raise FeatureNotFoundError()

    def verify_after_init(self):
        pass

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            for group in [PamGroup.AUTH, PamGroup.ACCOUNT]:
                if len(Fc.getPamGroup(self._pamService._data, group)) > 0:
                    raise ModifyError("group \"%s\" already has rule(s)" % (group), self._pamService.file_path)
            for group in [PamGroup.AUTH, PamGroup.ACCOUNT]:
                Fc.putPamGroup(target_data, group).append(PamRule(PamControl.REQUIRED, self.SO))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class Deny(_DenyOrPermit):

    SO = "pam_deny.so"

    @staticmethod
    def name():
        return "deny"


class Permit(_DenyOrPermit):

    SO = "pam_permit.so"

    @staticmethod
    def name():
        return "permit"


class _FallbackToPermitOrDeny(FeatureBase):

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        ruleListAuth, indexListAuth = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        ruleListAcct, indexListAcct = Fc.getPamGroup(self._pamService._data, PamGroup.ACCOUNT, lambda x: x.module_path == self.SO)
        if len(ruleListAuth) == 0 and len(ruleListAcct) == 0:
            raise FeatureNotFoundError()
        if (len(ruleListAuth) > 0 and len(ruleListAcct) == 0) or (len(ruleListAuth) == 0 and len(ruleListAcct) > 0):
            raise VerifyError("%s must be used for group \"%s\" and \"%s\" simultaneously" % (self.SO, PamGroup.AUTH, PamGroup.ACCOUNT), self._pamService.file_path)

        for pamGroup, ruleList, indexList in [(PamGroup.AUTH, ruleListAuth, indexListAuth), (PamGroup.ACCOUNT, ruleListAcct, indexListAcct)]:
            if len(ruleList) > 1:
                raise VerifyError("%s used for multiple times in group \"%s\"" % (self.SO, pamGroup), self._pamService.file_path)
            if ruleList[0].module_arguments is not None:
                raise VerifyError("%s in group \"%s\" should not have any argument" % (self.SO, pamGroup), self._pamService.file_path)
            if ruleList[0].control != PamControl.REQUIRED:
                raise VerifyError("%s in group \"%s\" can only be used for control \"%s\"" % (self.SO, pamGroup, PamControl.REQUIRED), self._pamService.file_path)
            if indexList[0] != len(self._pamService._data[pamGroup]) - 1:
                raise VerifyError("%s must be the last rule in group \"%s\"" % (self.SO, pamGroup), self._pamService.file_path)

    def modify(self, target_data, add_or_remove):
        if add_or_remove:
            if Fc.groupsHasModule(target_data, "*", self.SO):
                raise ModifyError("%s exists" % (self.SO), self._pamService.file_path)
            Fc.putPamGroup(target_data, PamGroup.AUTH).append(PamRule(PamControl.REQUIRED, self.SO))
            Fc.putPamGroup(target_data, PamGroup.ACCOUNT).append(PamRule(PamControl.REQUIRED, self.SO))
        else:
            Fc.pamGroupRemoveModule(target_data, "*", self.SO)


class FallbackToDeny(_FallbackToPermitOrDeny):

    SO = "pam_deny.so"

    @staticmethod
    def name():
        return "fallback-to-deny"

    def verify_after_init(self):
        pass


class FallbackToPermit(_FallbackToPermitOrDeny):

    SO = "pam_permit.so"

    @staticmethod
    def name():
        return "fallback-to-permit"

    def verify_after_init(self):
        raise VerifyError("fallback to permit is not allowed", self._pamService.file_path)


class FailLock(_FallbackToPermitOrDeny):

    SO = "pam_faillock.so"

    def __init__(self, context, pam_service):
        super().__init__(context, pam_service)

        self._confFile = None
        self._tallyDir = None

        ruleListAuth, _ = Fc.getPamGroup(self._pamService._data, PamGroup.AUTH, lambda x: x.module_path == self.SO)
        ruleListAcct, _ = Fc.getPamGroup(self._pamService._data, PamGroup.ACCOUNT, lambda x: x.module_path == self.SO)
        if len(ruleListAuth) == 0 and len(ruleListAcct) == 0:
            raise FeatureNotFoundError()
        if (len(ruleListAuth) > 0 and len(ruleListAcct) == 0) or (len(ruleListAuth) == 0 and len(ruleListAcct) > 0):
            raise VerifyError("%s must be used for group \"%s\" and \"%s\" simultaneously" % (self.SO, PamGroup.AUTH, PamGroup.ACCOUNT), self._pamService.file_path)

        # auth
        if len(ruleListAuth) in [2, 3]:
            if len(ruleListAuth[0].module_arguments) == 0 or ruleListAuth[0].module_arguments[0] != "preauth":
                raise VerifyError("the 1st appearance of %s must be \"preauth\" in group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)
            if len(ruleListAuth[1].module_arguments) == 0 or ruleListAuth[1].module_arguments[0] != "authfail":
                raise VerifyError("the 2nd appearance of %s must be \"authfail\" in group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

            if len(ruleListAuth) == 3:
                if len(ruleListAuth[2].module_arguments) == 0 or ruleListAuth[2].module_arguments[0] != "authsucc":
                    raise VerifyError("the 3rd appearance of %s must be \"authsucc\" in group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)

            for i, item in enumerate(ruleListAuth):
                for arg in item.module_arguments:
                    if arg.startswith("conf="):
                        if i == 0:
                            self._confFile = arg[len("conf="):]
                        else:
                            if arg[len("conf="):] != self._confFile:
                                raise VerifyError("all appearance of %s must have same \"conf=\" arguments in group \"%s\"" % (self.SO, PamGroup.AUTH), self._pamService.file_path)
                        break

        # account
        if True:
            if len(ruleListAcct) > 1:
                raise VerifyError("%s used for multiple times in group \"%s\"" % (self.SO, PamGroup.ACCOUNT), self._pamService.file_path)
            if ruleListAcct[0].module_arguments is not None:
                raise VerifyError("%s in group \"%s\" should not have any argument" % (self.SO, PamGroup.ACCOUNT), self._pamService.file_path)
            if ruleListAcct[0].control != PamControl.REQUIRED:
                raise VerifyError("%s in group \"%s\" can only be used for control \"%s\"" % (self.SO, PamGroup.ACCOUNT, PamControl.REQUIRED), self._pamService.file_path)

        # auth & account
        for i, item in enumerate(ruleListAuth + ruleListAcct):
            for arg in item.module_arguments:
                if arg.startswith("dir="):
                    if i == 0:
                        self._tallyDir = arg[len("dir="):]
                    else:
                        if arg[len("dir="):] != self._tallyDir:
                            raise VerifyError("all appearance of %s must have same \"dir=\" arguments" % (self.SO), self._pamService.file_path)
                    break

    @staticmethod
    def name():
        return "faillock"

    def verify_after_init(self):
        pass

    def modify(self, target_data, add_or_remove):
        assert False

    def is_using_conf_file(self):
        return self._confFile != "/dev/null"

    def is_using_default_conf_file(self):
        return self._confFile is None

    def get_conf_file(self):
        if self._confFile is None:
            return "/etc/security/faillock.conf"
        else:
            return self._confFile

    def is_using_default_tally_dir(self):
        return self._tallyDir is None

    def get_tally_dir(self):
        if self._tallyDir is None:
            return "/var/run/faillock"
        else:
            return self._tallyDir


__all__ = [
    "RootOk",
    "WheelOk",
    "NoLogin",
    "Shells",
    "UnixLogin",
    "Permit",
    "Deny",
    "FallbackToDeny",
    "FallbackToPermit",
    "MustBeWheel",
    "FailLock",
]
